package com.fanshuai.tree.bst;

import java.util.ArrayList;
import java.util.List;

/**
 * AVL平衡树
 * 节点左子树和右子树高度差小于2
 * 节点左子树也右子树也为AVL平衡树
 *
 * 平衡因子BF: 左子树和右子树高度差，AVL平衡树高度差为-1，0，1
 * 为了达到平衡，在插入和删除节点时，需要旋转子树。旋转分为左旋、右旋。旋转使得子树平衡，且不破坏二叉排序
 * @param <K>
 * @param <V>
 */
public class AVLTree<K extends Comparable<K>, V> implements BST<K, V> {
    class Node {
        public K key;
        public V value;
        public Node left;
        public Node right;

        public Node(K k, V v) {
            this.key = k;
            this.value = v;
            this.left = null;
            this.right = null;
        }
    }

    private Node root;
    private int size;
    public AVLTree() {
        root = null;
        size = 0;
    }

    /**
     * 节点高度
     * @param node
     * @return
     */
    private int getHeight(Node node) {
        if (null == node) {
            return 0;
        }
        return Math.max(getHeight(node.left), getHeight(node.right)) + 1;
    }

    /**
     * 平衡因子
     * @return
     */
    private int BF(Node node) {
        if (null == node) {
            return 0;
        }
        return getHeight(node.left) - getHeight(node.right);
    }

    /**
     * 左旋：右子树高度大于左子树，需要左旋。
     * 将右孩子节点设为父节点，父节点设为右孩子左节点。父节点右孩子设为右孩子左子树
     * @param node
     * @return
     */
    private Node rotateLeft(Node node) {
        //旋转树
        Node x = node.right;
        node.right = x.left;
        x.left = node;

        //返回旋转后的节点
        return x;
    }

    /**
     * 右旋：左子树高度大于右子树，需要右旋
     * 将父节点左孩子设置父节点，父节点设为左孩子右节点。父节点左孩子设为左孩子右子树
     * @param node
     * @return
     */
    private Node rotateRight(Node node) {
        Node x = node.left;
        node.left = x.right;
        x.right = node;
        return x;
    }

    /**
     * 按二叉排序树递归插入，回溯时旋转树，保持平衡
     * @param key
     * @param value
     */
    @Override
    public void put(K key, V value) {
        root = put(root, key, value);
    }

    private Node put(Node node, K key, V value) {
        if (null == node) {
            size++;
            return new Node(key, value);
        }

        if (key.compareTo(node.key) < 0) {
            node.left = put(node.left, key, value);
        } else if (key.compareTo(node.key) > 0) {
            node.right = put(node.right, key, value);
        } else {
            node.value = value;
        }

        return balance(node);
    }

    /**
     * 旋转树，保持平衡
     * 四种情况：LL，RR, LR, RL
     */
    private Node balance(Node node) {
        int bf = BF(node);
        if (bf >= 2) {
            int leftBF = BF(node.left);
            if (leftBF < 0) {
                //LR
                node = LRCase(node);
            } else {
                node = LLCase(node);
            }
        }
        if (bf <= -2) {
            int rightBF = BF(node.right);
            if (rightBF > 0) {
                node = RLCase(node);
            } else {
                node = RRCase(node);
            }
        }

        return node;
    }

    private Node LLCase(Node node) {
        return rotateRight(node);
    }

    private Node LRCase(Node node) {
        node.left = rotateLeft(node.left);
        return rotateRight(node);
    }

    private Node RRCase(Node node) {
        return rotateLeft(node);
    }

    private Node RLCase(Node node) {
        node.right = rotateRight(node.right);
        return rotateLeft(node);
    }

    @Override
    public V get(K key) {
        if (null == root) {
            return null;
        }

        Node node = root;
        while (node != null) {
            if (key.compareTo(node.key) == 0) {
                return node.value;
            } else if (key.compareTo(node.key) > 0) {
                node = node.right;
            } else {
                node = node.left;
            }
        }

        return null;
    }

    @Override
    public void delete(K key) {
        root = delete(root, key);
    }

    private Node delete(Node node, K key) {
        if (null == node) {
            return null;
        }

        boolean hit = false;
        if (key.compareTo(node.key) < 0) {
            node.left = delete(node.left, key);
        } else if (key.compareTo(node.key) > 0) {
            node.right = delete(node.right, key);
        } else {
            //返回节点
            Node newNode = null;

            size--;
            hit = true;
            if (null == node.left) {
                newNode = node.right;
                node.right = null;
                node = newNode;
            } else if (null == node.right) {
                newNode = node.left;
                node.left = null;
                node = newNode;
            } else {
                Node next = findMin(node.right);
                node.key = next.key;
                node.value = next.value;

                node.right = delete(node.right, next.key);
            }
        }

        /**
         * 旋转树，保持平衡
         */
        if (hit) {
            node = balance(node);
        }

        return node;
    }

    private Node findMin(Node node) {
        Node t = node;
        while (t.left != null) {
            t = t.left;
        }

        return t;
    }

    private Node findMax(Node node) {
        Node t = node;
        while (t.right != null) {
            t = t.right;
        }

        return t;
    }

    private void swapValue(Node p, Node q) {
        K k = p.key;
        V v = p.value;

        p.key = q.key;
        p.value = q.value;

        q.key = k;
        q.value = v;
    }

    @Override
    public List<K> keys() {
        List<K> keys = new ArrayList<>();
        fillKeys(root, keys);
        return keys;
    }

    private void fillKeys(Node node, List<K> keys) {
        if (null == node) {
            return;
        }

        fillKeys(node.left, keys);
        keys.add(node.key);
        fillKeys(node.right, keys);
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean empty() {
        return size == 0;
    }
}
